from utils import get_db_connection, generate_new_table_name, get_or_create_mapping_id
import psycopg2.sql as sql

class MappingTableOperations:
    def __init__(self, config):
        self.config = config
        self.connection = get_db_connection(config)

    def create_sensitive_data_table(self, table_name, columns):
        with self.connection.cursor() as cursor:
            if self.config.db_type == "postgresql":
                column_definitions = ", ".join([f"{col} TEXT" for col in columns])
                cursor.execute(sql.SQL("""
                    CREATE TABLE IF NOT EXISTS {} (
                        id SERIAL PRIMARY KEY,
                        {}
                    )
                """).format(sql.Identifier(table_name), sql.SQL(column_definitions)))
            elif self.config.db_type == "mysql":
                column_definitions = ", ".join([f"{col} TEXT" for col in columns])
                cursor.execute(sql.SQL("""
                    CREATE TABLE IF NOT EXISTS {} (
                        id INT AUTO_INCREMENT PRIMARY KEY,
                        {}
                    )
                """).format(sql.Identifier(table_name), sql.SQL(column_definitions)))
        self.connection.commit()

    def insert_sensitive_data(self, table_name, sensitive_columns, new_table_name):
        with self.connection.cursor() as cursor:
            for column in sensitive_columns:
                if self.config.db_type == "postgresql":
                    cursor.execute(sql.SQL("""
                        SELECT DISTINCT {} FROM {} WHERE {} IS NOT NULL
                    """).format(sql.Identifier(column), sql.Identifier(table_name), sql.Identifier(column)))
                elif self.config.db_type == "mysql":
                    cursor.execute(sql.SQL("""
                        SELECT DISTINCT {} FROM {} WHERE {} IS NOT NULL
                    """).format(sql.Identifier(column), sql.Identifier(table_name), sql.Identifier(column)))
                
                for row in cursor.fetchall():
                    value = row[0]
                    mapping_id = get_or_create_mapping_id(cursor, new_table_name, column, value)
                    self.connection.commit()

    def close_connection(self):
        self.connection.close()